import logging
import random

from flask import Blueprint
from flask import jsonify
from flask import request
from datetime import datetime, timedelta
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm.session import Session

from airflow.patsnap.service import workflow_service, backfill_service, alert_service
from airflow.configuration import conf
from airflow.models.dagrun import DagRun
from airflow.models.idata_workflow import WorkFlow
from airflow.utils import dates
from airflow.utils.session import provide_session
from airflow.utils.state import State
from airflow.www.api import exception_handle

log = logging.getLogger(__name__)

workflow_bp = Blueprint('workflow_bp', __name__)


@workflow_bp.route('/workflows/<int:workflow_id>', methods=['GET'])
@exception_handle
def get_workflow(workflow_id):
    workflow = WorkFlow.get_workflow(workflow_id)
    if workflow:
        return jsonify(code=0, message='ok', data=workflow.to_json())
    else:
        return jsonify(code=404, message='task not found', data=None)


@workflow_bp.route('/workflows', methods=['GET'])
@exception_handle
def query_workflow():
    project_id = request.args.get('project_id', None)
    folder_id = request.args.get('folder_id', None)
    if project_id is None and folder_id is None:
        return jsonify(code=1, message='project_id or folder_id 不能为空', data=None)
    name = request.args.get('name', None)
    tasks = WorkFlow.query_workflow(project_id=project_id, folder_id=folder_id, name=name)
    rs = [task.to_json() for task in tasks]
    return jsonify(code=0, message='ok', data=rs)


@workflow_bp.route('/workflows/<int:workflow_id>/history', methods=['GET'])
@exception_handle
def workflow_history(workflow_id):
    dag_id = 'dag-{}'.format(workflow_id)
    return jsonify(code=0, message='ok', data=workflow_service.workflow_history(dag_id))


@workflow_bp.route('/workflows/instances', methods=['GET'])
@exception_handle
def workflow_instances():
    page = request.args.get('page', '1')
    page = int(page)
    page_size = request.args.get('page_size', '10')
    page_size = int(page_size)
    dag_id = request.args.get('dag_id', None)
    project_id = request.args.get('project_id', None)
    if not project_id:
        return jsonify(code=400, message='project_id 不能为空', data=None)
    dag_name = request.args.get('dag_name', None)
    state = request.args.get('state', None)
    run_type = request.args.get('run_type', None)
    execution_date_from = request.args.get('execution_date_from', None)
    if execution_date_from:
        execution_date_from = int(execution_date_from)
    execution_date_to = request.args.get('execution_date_to', None)
    if execution_date_to:
        execution_date_to = int(execution_date_to)

    return jsonify(code=0, message='ok',
                   data=workflow_service.workflow_instances(project_id=project_id, page=page,
                                                            page_size=page_size,
                                                            dag_id=dag_id, dag_name=dag_name,
                                                            state=state, run_type=run_type,
                                                            execution_date_from=execution_date_from,
                                                            execution_date_to=execution_date_to))


@workflow_bp.route('/workflows/run_types', methods=['GET'])
@exception_handle
def get_dag_run_types():
    data = [
        {
            'label': 'SCHEDULED',
            'value': 'scheduled'
        },
        {
            'label': 'MANUAL',
            'value': 'manual'
        },
        {
            'label': 'BACKFILL_JOB',
            'value': 'backfill'
        }
    ]
    return jsonify(code=0, message='ok', data=data)


@workflow_bp.route('/workflows/scheduler_interval', methods=['GET'])
@exception_handle
def get_scheduler_interval():
    return jsonify(code=0, message='ok', data=workflow_service.scheduler_interval)


@workflow_bp.route('/workflows/cron', methods=['GET'])
@exception_handle
def get_default_cron():
    interval = request.args.get('interval', None)
    if interval:
        return jsonify(code=0, message='ok', data=workflow_service.get_default_cron(interval))
    else:
        return jsonify(code=400, message='bad request interval not found', data=None)


@workflow_bp.route('/workflows', methods=['POST'])
@exception_handle
def insert():
    data = request.get_json()
    workflow = WorkFlow.from_dict(data)
    try:
        workflow.insert()
    except IntegrityError as err:
        return jsonify(code=500, message=str(err), data=None)
    return jsonify(code=0, message='ok', data=workflow.to_json())


@workflow_bp.route('/workflows/<int:workflow_id>', methods=['PUT'])
@exception_handle
def update(workflow_id):
    data = request.get_json()
    if not 'id' in data:
        data['id'] = workflow_id
    workflow = WorkFlow.from_dict(data)
    workflow.update()
    return jsonify(code=0, message='ok', data=workflow.to_json())


@workflow_bp.route('/workflows/<int:workflow_id>', methods=['PATCH'])
@exception_handle
def patch(workflow_id):
    data = request.get_json()
    dag_id = 'dag-{}'.format(workflow_id)
    running_workflow_instance_count = workflow_service.workflow_instance_count(dag_id, 'running')
    if running_workflow_instance_count > 0:
        error = '当前workflow正在运行中，无法修改，请稍后运行完再修改。'
        return jsonify(code=400, message=error, data=None)

    error = workflow_service.workflow_validate(workflow_id, data)
    if error is None:
        workflow_service.workflow_patch(workflow_id, data)
        return jsonify(code=0, message='ok', data=None)
    else:
        return jsonify(code=400, message=error, data=None)


@workflow_bp.route('/workflows/<int:workflow_id>', methods=['DELETE'])
@exception_handle
def delete(workflow_id):
    wk = WorkFlow.get_workflow(workflow_id)
    if not wk.is_paused:
        return jsonify(code=400, message='开启定时调度的工作流不能删除', data=None)
    WorkFlow.delete(workflow_id)
    return jsonify(code=0, message='ok', data=None)


@workflow_bp.route('/workflows/<int:workflow_id>/backfill', methods=['POST'])
@exception_handle
def back_fill(workflow_id):
    data = request.get_json()
    start_date = data.get('start_date', None)
    end_date = data.get('end_date', None)
    is_synchronous = data.get('is_synchronous',  False)
    synchronous_num = data.get('synchronous_num',  1)
    if None in [start_date, end_date]:
        return jsonify(code=400, message='start_date or end_date is None', data=None)
    start_date = dates.from_milliseconds(start_date)
    end_date = dates.from_milliseconds(end_date)
    dag_id = 'dag-{}'.format(workflow_id)
    backfill_service.check(dag_id)
    workflow_service.clear_dag_run(dag_id, start_date, end_date)
    backfill_service.run_in_pod(dag_id, start_date, end_date, is_synchronous, synchronous_num)
    return jsonify(code=0, message='ok', data=None)


@workflow_bp.route('/backfill/<int:job_id>/kill', methods=['POST'])
@exception_handle
def back_kill(job_id):
    backfill_service.kill_pod(job_id)
    backfill_service.set_backfill_job_failed(job_id)
    return jsonify(code=0, message='ok', data=None)


@workflow_bp.route('/workflows/<int:workflow_id>/backfill', methods=['GET'])
@exception_handle
def back_query(workflow_id):
    dag_id = 'dag-{}'.format(workflow_id)
    return jsonify(code=0, message='ok', data=backfill_service.query_job(dag_id))


@workflow_bp.route('/workflows/<int:workflow_id>/trigger', methods=['POST'])
@exception_handle
@provide_session
def workflow_trigger(workflow_id, session: Session = None):
    import time
    execution_date = None
    data = request.get_json()
    if data:
        execution_date = data.get('execution_date', None)
    if not execution_date:
        return jsonify(code=400, message='计划时间不能为空', data=None)
    wk = WorkFlow.get_workflow(workflow_id, session=session)
    if execution_date < dates.to_milliseconds(wk.start_time):
        return jsonify(code=400, message='计划时间不能小于工作流开始时间', data=None)
    if execution_date > int(time.time() * 1000):
        return jsonify(code=400, message='计划时间不能选择未来时间', data=None)
    dag_id = 'dag-{}'.format(workflow_id)
    max_runs = conf.getint('core', 'max_active_runs_per_dag')
    rs = DagRun.find(dag_id=dag_id, state=State.RUNNING, session=session)
    if len(rs) >= max_runs:
        return jsonify(code=1, message='超过单个工作流最大并行运行数 {}'.format(max_runs), data=None)
    if execution_date:
        execution_date = dates.from_milliseconds(execution_date + random.randint(0, 1000))
    else:
        execution_date = dates.sh_now()
    workflow_service.create_dag_run(dag_id, execution_date, session=session)
    return jsonify(code=0, message='ok',
                   data={'dag_id': dag_id, 'execution_date': dates.to_milliseconds(execution_date)})


@workflow_bp.route('/workflows/<int:workflow_id>/children', methods=['GET'])
@exception_handle
@provide_session
def workflow_children(workflow_id, session: Session = None):
    wks = session.query(WorkFlow).filter(WorkFlow.parent_id == workflow_id).all()
    rs = [it.to_json() for it in wks]
    return jsonify(code=0, message='ok', data=rs)


@workflow_bp.route('/workflows/kill', methods=['POST'])
@exception_handle
def workflow_kill():
    data = request.get_json()
    dag_id = data.get('dag_id', None)
    execution_date = data.get('execution_date', None)
    if None in [dag_id, execution_date]:
        return jsonify(code=400, message='dag_id or execution_date is None', data=None)
    execution_date = dates.from_milliseconds(execution_date)
    workflow_service.mark_dag_failed(dag_id, execution_date)
    return jsonify(code=0, message='ok', data=None)


@workflow_bp.route('/workflows/options', methods=['GET'])
@exception_handle
def workflow_options():
    project_id = request.args.get('project_id', None)
    if not project_id:
        return jsonify(code=400, message='project_id not found', data=None)
    rs = workflow_service.query_options_by_project(int(project_id))
    return jsonify(code=0, message='ok', data=rs)


@workflow_bp.route('/workflows/state_agg', methods=['GET'])
@exception_handle
def workflow_state_agg():
    start_time = request.args.get('start_time', None, type=int)
    end_time = request.args.get('end_time', None, type=int)
    project_id = request.args.get('project_id', None, type=int)
    if start_time is None:
        log.error("workflow_state_agg error start_time=%s", start_time, )
        return jsonify(code=400, message="task_state_agg error start_time=None", data=None)
    if end_time is None:
        log.error("workflow_state_agg error end_time=%s", end_time, )
        return jsonify(code=400, message="task_state_agg error end_time=None", data=None)

    rs = workflow_service.get_by_start_date(project_id, start_time, end_time)
    return jsonify(code=0, message='ok', data=rs)


@workflow_bp.route('/dagrun/diagnose', methods=['GET'])
@exception_handle
def diagnose_dagrun():
    dag_id = request.args.get('dag_id', None)
    execution_date = request.args.get('execution_date', None, type=int)
    if dag_id is None or execution_date is None:
        return jsonify(code=400, message="dag_id or execution_date is None", data=None)
    execution_date = dates.from_milliseconds(execution_date)
    rs = workflow_service.diagnose_dagrun(dag_id, execution_date)
    return jsonify(code=0, message='ok', data=rs)


@workflow_bp.route('/workflows/<int:workflow_id>/alerts', methods=['POST'])
@exception_handle
def delete_workflow_alert(workflow_id):
    condition = int(request.args.get('condition', 0))
    alerts = request.get_json()
    alert_service.save_workflow_alerts(workflow_id, alerts, condition=condition)
    return jsonify(code=0, message='ok', data=None)


@workflow_bp.route('/workflows/<int:workflow_id>/alerts', methods=['GET'])
@exception_handle
def get_workflow_alert(workflow_id):
    condition = int(request.args.get('condition', 0))
    rs = alert_service.get_workflow_alerts_display(workflow_id, condition=condition)
    return jsonify(code=0, message='ok', data=rs)


@workflow_bp.route('/workflows/clear_history', methods=['POST'])
@exception_handle
def clear_history():
    today = datetime.today()
    date_before = today - timedelta(days=90)
    data = request.get_json()
    dag_id = data.get('dag_id', None)
    start_date = data.get('start_date', None)
    end_date = data.get('end_date', None)
    if dag_id and start_date and end_date:
        start_date = dates.from_milliseconds(start_date)
        end_date = dates.from_milliseconds(end_date)
        workflow_service.clear_dag_run(dag_id, start_date, end_date)
    else:
        workflow_service.clear_history_not_running(date_before)
    return jsonify(code=0, message='ok', data=None)


@workflow_bp.route('/workflows/retry_failed', methods=['POST'])
@exception_handle
def workflow_retry_failed():
    data = request.get_json()
    dag_id = data.get('dag_id', None)
    execution_date = data.get('execution_date', None)
    if dag_id and execution_date:
        date = dates.from_milliseconds(execution_date)
        rs = workflow_service.retry_failed(dag_id=dag_id, execution_date=date)
        return jsonify(code=0, message='ok', data=rs)
    else:
        return jsonify(code=400, message='dag_id or execution_date 不能为空', data=None)
